package de.unibi.comet.fa;

import java.util.Arrays;
import java.util.HashMap;
import java.util.Iterator;
import java.util.LinkedList;
import java.util.List;
import java.util.NoSuchElementException;

/** Represents a partition of the integers 0,..,n into blocks (or subsets). */
public class Partition {
	// data for each element
	/** Index of next element in same block. */
	private int[] next;
	/** Index of last element in same block. */
	private int[] last;
	/** For each integer the block it belongs to. */
	private int[] block;
	
	// data for each block
	/** Size of each block. */
	private int[] blockSize;
	/** For each block the index of the first element in the block. */
	private int[] firstElement;
	/** List if empty blocks. */
	private LinkedList<Integer> emptyBlocks;
	
	private int blockCount;
		
	/** Construct a partition of given number of integers. At first, all integers
	 *  are in the same block. */
	public Partition(int size) {
		next = new int[size];
		last = new int[size];
		for (int i=0; i<size-1; ++i) {
			next[i]=i+1;
			last[i]=i-1;
		}
		next[size-1]=-1;
		last[size-1]=size-2;

		block = new int[size];

		blockSize = new int[size];
		blockSize[0]=size;
		
		firstElement = new int[size];
		Arrays.fill(firstElement, -1);
		firstElement[0]=0;
		
		emptyBlocks = new LinkedList<Integer>();
		for (int i=1; i<size; ++i) emptyBlocks.add(i);
		
		blockCount=1;
	}
	
	/** Construct a partition from an initial partition. Where i and j are put into the same
	 *  block, iff initialPartition[i].equals(initialPartition[j]).
	 */
	public Partition(List initialPartition) {
		int size = initialPartition.size();
		next = new int[size];
		Arrays.fill(next, -1);
		last = new int[size];
		Arrays.fill(last, -1);
		
		blockSize = new int[size];
		
		firstElement = new int[size];
		Arrays.fill(firstElement, -1);
		
		int[] lastInBlock = new int[size];
		Arrays.fill(lastInBlock, -1);

		// maps numbers in array initialPartition to block indices
		HashMap<Object, Integer> blockIndexMap = new HashMap<Object, Integer>(2*size);
		blockCount=0;
		block = new int[size];
		int i = 0;
		for (Object j : initialPartition) {
			Integer k = blockIndexMap.get(j);
			if (k==null) {
				k=blockIndexMap.size();
				blockIndexMap.put(j, k);
				++blockCount;
			}
			block[i]=k;
			if (firstElement[k]==-1) {
				firstElement[k]=i; 
			} else {
				next[lastInBlock[k]]=i;
			}
			last[i]=lastInBlock[k];
			blockSize[k]+=1;
			lastInBlock[k]=i;
			++i;
		}
		
		emptyBlocks = new LinkedList<Integer>();	
		for (int j=0; j<size; ++j) {
			if (blockSize[j]==0) emptyBlocks.add(j);
		}
	}
	
	private class BlockIterator implements Iterator<Integer> {
		private int nextElement;
		BlockIterator(int blockIndex) { 
			this.nextElement=firstElement[blockIndex];
		}
		public boolean hasNext() { return nextElement>=0; }
		public Integer next() {
			if (nextElement<0) throw new NoSuchElementException();
			int result = nextElement;
			nextElement = next[nextElement];
			return result;
		}
		public void remove() { throw new UnsupportedOperationException(); }
	}
	
	private class IterableBlock implements Iterable<Integer> {
		private int blockIndex;
		IterableBlock (int blockIndex) {
			this.blockIndex=blockIndex;
		}
		public Iterator<Integer> iterator() { return new BlockIterator(blockIndex); } 
	}
	
	/** Returns an Iterable object that allow iteration over elements of a block. */
	public Iterable<Integer> block(int blockIndex) {
		if ((blockIndex<0)||(blockIndex>=blockSize.length)) throw new IndexOutOfBoundsException();
		return new IterableBlock(blockIndex);
	}
	
	/** Split the given integers off from their current block and let 
	 *  them form a new block. The number of the new block is returned
	 *  (runtime O(l.size())). */
	public int splitOff(List<Integer> l) {
		if (l.size()==0) throw new IllegalArgumentException("Cannot split off an empty block.");
		// if no block is free, set block index to -1 and assign a number later on,
		// when a block is empty
		int newBlockIndex = (emptyBlocks.size()>0)?emptyBlocks.remove():-1;
		++blockCount;
		int newBlockFirst = l.get(0);
		int newBlockLast = -1;
		int newBlockSize = 0;
		for (int i : l) {
			// ignore duplicates
			if (block[i]==newBlockIndex) continue;
			// take care of blocksizes
			int oldBlockIndex = block[i]; 
			blockSize[oldBlockIndex]-=1;
			if (blockSize[oldBlockIndex]==0) {
				emptyBlocks.add(oldBlockIndex);
				--blockCount;
			}
			block[i]=newBlockIndex;
			newBlockSize+=1;
			// unlink element from its old block
			if (firstElement[oldBlockIndex]==i) firstElement[oldBlockIndex]=next[i];
			if (last[i]!=-1) next[last[i]]=next[i];
			if (next[i]!=-1) last[next[i]]=last[i];
			// link element to its new block
			next[i]=-1;
			if (newBlockLast>=0) next[newBlockLast]=i;
			last[i]=newBlockLast;
			newBlockLast=i;
		}
		if (newBlockIndex==-1) {
			newBlockIndex=emptyBlocks.remove();
			int i=newBlockFirst;
			while (i!=-1) {
				block[i]=newBlockIndex;
				i=next[i];
			}
		}
		blockSize[newBlockIndex]=newBlockSize;
		firstElement[newBlockIndex]=newBlockFirst;
		return newBlockIndex;
	}
	
	/** Returns the index of the block the given element belongs to
	 *  (runtime O(1)). */
	public int getBlockIndex(int element) { 
		return block[element];
	}
	
	/** Returns the number of blocks. */
	public int blockCount() {
		return blockCount;
	}
	
	private class BlockListIterator implements Iterator<Integer> {
		private int nextElement;
		private int count;
		BlockListIterator() {
			for (int i=0; i<blockSize.length; ++i) {
				if (blockSize[i]>0) {
					nextElement=i;
					break;
				}
			}
			count=1;
		}
		public boolean hasNext() { return nextElement>=0; }
		public Integer next() {
			if (nextElement<0) throw new NoSuchElementException();
			int result = nextElement;
			if (count==blockCount) {
				nextElement=-1;
			} else {
				for (int i=nextElement+1; i<blockSize.length; ++i) {
					if (blockSize[i]>0) {
						nextElement=i;
						break;
					}
				}
				++count;
			}
			return result;
		}
		public void remove() { throw new UnsupportedOperationException(); }
	}
	
	private class IterableBlockList implements Iterable<Integer> {
		IterableBlockList () {}
		public Iterator<Integer> iterator() { return new BlockListIterator(); } 
	}

	public Iterable<Integer> blockIndices() {
		return new IterableBlockList();
	}
	
	public int blockSize(int blockIndex) {
		return blockSize[blockIndex];
	}

	public int size() {
		return block.length; 
	}
	
	@Override
	public String toString() {
		StringBuilder sb = new StringBuilder();
		for (int i=0; i<blockCount(); ++i) {
			sb.append(String.format("%d: ", i));
			for (int j : block(i)) {
				sb.append(String.format("%d, ", j));
			}
			sb.append("\n");
		}
		return sb.toString();
	}
	
}
